import pandas as pd
import os

print("数据预处理开始~")

# 如果当前路径没有"_1.截面数据"文件夹，则创建该文件夹
if not os.path.exists("_1.截面数据"):
    os.mkdir("_1.截面数据")

# 读取"1.月度涨跌幅.xlsx"、"2.市盈率.xlsx"、"3.净资产收益率ROE.xlsx"这三个文件
df_1 = pd.read_excel("源数据\\1.月度涨跌幅.xlsx")
df_2 = pd.read_excel("源数据\\2.市盈率.xlsx")
df_3 = pd.read_excel("源数据\\3.净资产收益率ROE.xlsx")

# 提取每个月的数据，合并得到截面数据
cols = df_1.columns[2:]
for col in cols:
    df_csd = df_1.iloc[:, :2]  # 前两列都是一样的
    df_csd = pd.concat(objs=[df_csd, df_1[col], df_2[col], df_3[col]], axis=1)  # 把每个月的数据列都合并到df_csd
    df_csd.columns = ["证券代码", "证券名称", "月度涨跌幅", "市盈率", "净资产收益率ROE"]  # 重命名列名
    df_csd.to_excel("_1.截面数据\\" + col + ".xlsx", index=False)

print("数据预处理完成！")
